package com.bruce.ai.alibaba.workflow.controller;

import java.util.HashMap;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.MediaType;
import java.util.Objects;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import com.alibaba.cloud.ai.graph.CompileConfig;
import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.OverAllState;
 
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.checkpoint.config.SaverConfig;
import com.alibaba.cloud.ai.graph.agent.ReactAgent;
import com.alibaba.cloud.ai.graph.checkpoint.savers.MemorySaver;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;

import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * 普通节点、智能体节点混合工作流控制器
 */
@RequestMapping("/hybrid")
@RestController
public class HybridWorkflowController {

    private CompiledGraph compiledGraph ;
    private final ReactAgent qaAgent;

    private static final Logger logger = LoggerFactory.getLogger(HybridWorkflowController.class);

    public HybridWorkflowController(@Qualifier("hybridWorkflow") StateGraph hybridGraph,
                                    @Qualifier("qaAgent") ReactAgent qaAgent) throws Exception{
        SaverConfig saverConfig = SaverConfig.builder().register(new MemorySaver())
                .build();
        this.compiledGraph = hybridGraph
                .compile(CompileConfig.builder().saverConfig(saverConfig).build());
        this.qaAgent = qaAgent;
    }

    /**
     * 展开工作流
     */
    private String buildEvent(String node, String threadId, String type, Object payload) {
        java.util.Map<String, Object> evt = new java.util.HashMap<>();
        evt.put("node", node);
        evt.put("thread", threadId);
        evt.put("type", type);
        evt.put("payload", payload);
        return com.alibaba.fastjson.JSON.toJSONString(evt);
    }

    private String extractStreamText(Object output) {
        if (output instanceof com.alibaba.cloud.ai.graph.streaming.StreamingOutput<?> streaming) {
            try {
                java.lang.reflect.Method delta = streaming.getClass().getMethod("delta");
                Object d = delta.invoke(streaming);
                String s = String.valueOf(d);
                if (s != null && !s.isEmpty()) return s;
            } catch (Throwable ignored) {}
            try {
                java.lang.reflect.Method content = streaming.getClass().getMethod("content");
                Object c = content.invoke(streaming);
                String s = String.valueOf(c);
                if (s != null && !s.isEmpty()) return s;
            } catch (Throwable ignored) {}
            try {
                java.lang.reflect.Method chunk = streaming.getClass().getMethod("chunk");
                Object ch = chunk.invoke(streaming);
                String s = String.valueOf(ch);
                if (s != null) return s;
            } catch (Throwable ignored) {}
        }
        return String.valueOf(output);
    }

    @RequestMapping(value ="/expand", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> expand(@RequestParam("input") String input) throws Exception{
        Map<String, Object> objectMap = new HashMap<>();
        objectMap.put("input", input);
        return Mono.defer(() -> Mono.justOrEmpty(compiledGraph.invoke(objectMap))
                .map(OverAllState::data)
                .map(data -> buildEvent("hybrid", "context7", "final", String.valueOf(data)))
                .map(s -> ServerSentEvent.builder(java.util.Objects.requireNonNull(s)).build())
                .onErrorResume(e -> Mono.just(ServerSentEvent.builder(buildEvent("error", "context7", "error", e.getMessage())).build())))
            .flux()
            .onErrorResume(e -> Flux.just(ServerSentEvent.builder(buildEvent("error", "context7", "error", e.getMessage())).build()))
            .doOnCancel(() -> logger.info("Client disconnected from stream"))
            .doOnError(e -> logger.error("Error occurred during streaming", e));
    }

    @RequestMapping(value ="/expandStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> expandStream(@RequestParam("input") String input,
                                                      @RequestParam(value = "thread_id", defaultValue = "context7", required = false) String threadId) throws GraphRunnerException, GraphStateException {
        String cleaned = Objects.requireNonNullElse(input, "").trim().toLowerCase();
        String head = buildEvent("preprocess", threadId, "start", cleaned);

        Flux<ServerSentEvent<String>> headFlux = Flux.just(ServerSentEvent.<String>builder(java.util.Objects.requireNonNull(head)).build());

        StringBuilder buf = new StringBuilder();
        Flux<ServerSentEvent<String>> bodyFlux = qaAgent.stream(cleaned)
                .flatMap(output -> {
                    String piece = extractStreamText(output);
                    if (!piece.isEmpty()) {
                        buf.append(piece);
                        String evt = buildEvent("qa_agent", threadId, "stream", piece);
                        return Flux.just(ServerSentEvent.<String>builder(java.util.Objects.requireNonNull(evt)).build());
                    }
                    return Flux.empty();
                });

        Flux<ServerSentEvent<String>> tailFlux = Flux.defer(() -> {
            String answer = buf.toString();
            boolean isValid = answer.length() > 50;
            java.util.Map<String, Object> payload = java.util.Map.of("qa_result", answer, "is_valid", isValid);
            String finalEvt = buildEvent("validate", threadId, "final", payload);
            return Flux.just(ServerSentEvent.<String>builder(java.util.Objects.requireNonNull(finalEvt)).build());
        });

        return Flux.concat(headFlux, bodyFlux, tailFlux)
            .onErrorResume(e -> Flux.just(ServerSentEvent.<String>builder(buildEvent("error", threadId, "error", e.getMessage())).build()))
            .doOnCancel(() -> logger.info("Client disconnected from stream"))
            .doOnError(e -> logger.error("Error occurred during streaming", e));
    }


}
